#!/usr/bin/python

# test.py: seeing how fast we can load GTFS.

# Copyright 2011 Portland Transport

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import csv

stopsIn = csv.reader(open('stops.txt'))
routesIn = csv.reader(open('routes.txt'))
tripsIn = csv.reader(open('trips.txt'))
stopTimesIn = csv.reader(open('stop_times.txt'))

# First line has column names
stopsCols = stopsIn.next()
routesCols = routesIn.next()
tripsCols = tripsIn.next()
stopTimesCols = stopTimesIn.next()

stops = {}
for stop in stopsIn:
    stopDict = dict(zip(stopsCols, stop))
    stopDict['routes'] = []
    stops[stopDict['stop_id']] = stopDict

del stopsIn

print '%s stops in GTFS' % len (stops)

# now run through the routes
routes = {}
for route in routesIn:
    routeDict = dict(zip(routesCols, route))
    routes[routeDict['route_id']] = routeDict

del routesIn

print '%s routes in GTFS' % len (routes)

# Build a mapping of trips to routes
tripMap = {}
for trip in tripsIn:
    tripDict = dict(zip(tripsCols, trip))
    tripMap[tripDict['trip_id']] = tripDict['route_id']

del tripsIn

print '%s trips in GTFS' % len (tripMap)

for st in stopTimesIn:
    stDict = dict(zip(stopTimesCols, st))
    stop = stops[stDict['stop_id']]
    route = routes[tripMap[stDict['trip_id']]]
    if not route in stop['routes']:
        stop['routes'].append(route)

del stopTimesIn

print 'Done.'
raw_input('Press enter when done')
